{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h1 align=\"center\">Non-Rigid Registration: Free Form Deformation</h1>\n",
    "\n",
    "This notebook illustrates the use of the Free Form Deformation (FFD) based non-rigid registration algorithm in SimpleITK.\n",
    "\n",
    "The data we work with is a 4D (3D+time) thoracic-abdominal CT, the Point-validated Pixel-based Breathing Thorax Model (POPI) model. This data consists of a set of temporal CT volumes, a set of masks segmenting each of the CTs to air/body/lung, and a set of corresponding points across the CT volumes. \n",
    "\n",
    "The POPI model is provided by the Léon Bérard Cancer Center & CREATIS Laboratory, Lyon, France. The relevant publication is:\n",
    "\n",
    "J. Vandemeulebroucke, D. Sarrut, P. Clarysse, \"The POPI-model, a point-validated pixel-based breathing thorax model\",\n",
    "Proc. XVth International Conference on the Use of Computers in Radiation Therapy (ICCR), Toronto, Canada, 2007.\n",
    "\n",
    "The POPI data, and additional 4D CT data sets with reference points are available from the CREATIS Laboratory <a href=\"http://www.creatis.insa-lyon.fr/rio/popi-model?action=show&redirect=popi\">here</a>. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "\n",
    "import SimpleITK as sitk\n",
    "import registration_utilities as ru\n",
    "import registration_callbacks as rc\n",
    "\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from ipywidgets import interact, fixed\n",
    "\n",
    "#utility method that either downloads data from the Girder repository or\n",
    "#if already downloaded returns the file name for reading from disk (cached data)\n",
    "%run update_path_to_download_script\n",
    "from downloaddata import fetch_data as fdata"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Utilities\n",
    "\n",
    "Load utilities that are specific to the POPI data, functions for loading ground truth data, display and the labels for masks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run popi_utilities_setup.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading Data\n",
    "\n",
    "Load all of the images, masks and point data into corresponding lists. If the data is not available locally it will be downloaded from the original remote repository. \n",
    "\n",
    "Take a look at the images. According to the documentation on the POPI site, volume number one corresponds to end inspiration (maximal air volume)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images = []\n",
    "masks = []\n",
    "points = []\n",
    "for i in range(0,10):\n",
    "    image_file_name = 'POPI/meta/{0}0-P.mhd'.format(i)\n",
    "    mask_file_name = 'POPI/masks/{0}0-air-body-lungs.mhd'.format(i)\n",
    "    points_file_name = 'POPI/landmarks/{0}0-Landmarks.pts'.format(i)\n",
    "    images.append(sitk.ReadImage(fdata(image_file_name), sitk.sitkFloat32)) #read and cast to format required for registration\n",
    "    masks.append(sitk.ReadImage(fdata(mask_file_name)))\n",
    "    points.append(read_POPI_points(fdata(points_file_name)))\n",
    "        \n",
    "interact(display_coronal_with_overlay, temporal_slice=(0,len(images)-1), \n",
    "         coronal_slice = (0, images[0].GetSize()[1]-1), \n",
    "         images = fixed(images), masks = fixed(masks), \n",
    "         label=fixed(lung_label), window_min = fixed(-1024), window_max=fixed(976));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Getting to know your data\n",
    "\n",
    "While the POPI site states that image number 1 is end inspiration, and visual inspection seems to suggest this is correct, we should probably take a look at the lung volumes to ensure that what we expect is indeed what is happening.\n",
    "\n",
    "Which image is end inspiration and which end expiration?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "label_shape_statistics_filter = sitk.LabelShapeStatisticsImageFilter()\n",
    "\n",
    "for i, mask in enumerate(masks):\n",
    "    label_shape_statistics_filter.Execute(mask)\n",
    "    print('Lung volume in image {0} is {1} liters.'.format(i,0.000001*label_shape_statistics_filter.GetPhysicalSize(lung_label)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Free Form Deformation\n",
    "\n",
    "This function will align the fixed and moving images using a FFD. If given a mask, the similarity metric will be evaluated using points sampled inside the mask. If given fixed and moving points the similarity metric value and the target registration errors will be displayed during registration. \n",
    "\n",
    "As this notebook performs intra-modal registration, we use the MeanSquares similarity metric (simple to compute and appropriate for the task)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def bspline_intra_modal_registration(fixed_image, moving_image, fixed_image_mask=None, fixed_points=None, moving_points=None):\n",
    "\n",
    "    registration_method = sitk.ImageRegistrationMethod()\n",
    "    \n",
    "    # Determine the number of BSpline control points using the physical spacing we want for the control grid. \n",
    "    grid_physical_spacing = [50.0, 50.0, 50.0] # A control point every 50mm\n",
    "    image_physical_size = [size*spacing for size,spacing in zip(fixed_image.GetSize(), fixed_image.GetSpacing())]\n",
    "    mesh_size = [int(image_size/grid_spacing + 0.5) \\\n",
    "                 for image_size,grid_spacing in zip(image_physical_size,grid_physical_spacing)]\n",
    "\n",
    "    initial_transform = sitk.BSplineTransformInitializer(image1 = fixed_image, \n",
    "                                                         transformDomainMeshSize = mesh_size, order=3)    \n",
    "    registration_method.SetInitialTransform(initial_transform)\n",
    "        \n",
    "    registration_method.SetMetricAsMeanSquares()\n",
    "    # Settings for metric sampling, usage of a mask is optional. When given a mask the sample points will be \n",
    "    # generated inside that region. Also, this implicitly speeds things up as the mask is smaller than the\n",
    "    # whole image.\n",
    "    registration_method.SetMetricSamplingStrategy(registration_method.RANDOM)\n",
    "    registration_method.SetMetricSamplingPercentage(0.01)\n",
    "    if fixed_image_mask:\n",
    "        registration_method.SetMetricFixedMask(fixed_image_mask)\n",
    "    \n",
    "    # Multi-resolution framework.            \n",
    "    registration_method.SetShrinkFactorsPerLevel(shrinkFactors = [4,2,1])\n",
    "    registration_method.SetSmoothingSigmasPerLevel(smoothingSigmas=[2,1,0])\n",
    "    registration_method.SmoothingSigmasAreSpecifiedInPhysicalUnitsOn()\n",
    "\n",
    "    registration_method.SetInterpolator(sitk.sitkLinear)\n",
    "    registration_method.SetOptimizerAsLBFGSB(gradientConvergenceTolerance=1e-5, numberOfIterations=100)\n",
    "    \n",
    "\n",
    "    # If corresponding points in the fixed and moving image are given then we display the similarity metric\n",
    "    # and the TRE during the registration.\n",
    "    if fixed_points and moving_points:\n",
    "        registration_method.AddCommand(sitk.sitkStartEvent, rc.metric_and_reference_start_plot)\n",
    "        registration_method.AddCommand(sitk.sitkEndEvent, rc.metric_and_reference_end_plot)\n",
    "        registration_method.AddCommand(sitk.sitkIterationEvent, lambda: rc.metric_and_reference_plot_values(registration_method, fixed_points, moving_points))\n",
    "    \n",
    "    return registration_method.Execute(fixed_image, moving_image)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Perform Registration\n",
    "\n",
    "The following cell allows you to select the images used for registration, runs the registration, and afterwards computes statistics comparing the target registration errors before and after registration and displays a histogram of the TREs.\n",
    "\n",
    "To time the registration, uncomment the timeit magic. \n",
    "<b>Note</b>: this creates a separate scope for the cell. Variables set inside the cell, specifically tx, will become local variables and thus their value is not available in other cells. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#%%timeit -r1 -n1\n",
    "\n",
    "# Select the fixed and moving images, valid entries are in [0,9].\n",
    "fixed_image_index = 0\n",
    "moving_image_index = 7\n",
    "\n",
    "\n",
    "tx = bspline_intra_modal_registration(fixed_image = images[fixed_image_index], \n",
    "                                      moving_image = images[moving_image_index],\n",
    "                                      fixed_image_mask = (masks[fixed_image_index] == lung_label),\n",
    "                                      fixed_points = points[fixed_image_index], \n",
    "                                      moving_points = points[moving_image_index]\n",
    "                                     )\n",
    "initial_errors_mean, initial_errors_std, _, initial_errors_max, initial_errors = ru.registration_errors(sitk.Euler3DTransform(), points[fixed_image_index], points[moving_image_index])\n",
    "final_errors_mean, final_errors_std, _, final_errors_max, final_errors = ru.registration_errors(tx, points[fixed_image_index], points[moving_image_index])\n",
    "\n",
    "plt.hist(initial_errors, bins=20, alpha=0.5, label='before registration', color='blue')\n",
    "plt.hist(final_errors, bins=20, alpha=0.5, label='after registration', color='green')\n",
    "plt.legend()\n",
    "plt.title('TRE histogram');\n",
    "print('Initial alignment errors in millimeters, mean(std): {:.2f}({:.2f}), max: {:.2f}'.format(initial_errors_mean, initial_errors_std, initial_errors_max))\n",
    "print('Final alignment errors in millimeters, mean(std): {:.2f}({:.2f}), max: {:.2f}'.format(final_errors_mean, final_errors_std, final_errors_max))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another option for evaluating the registration is to use segmentation. In this case, we transfer the segmentation from one image to the other and compare the overlaps, both visually, and quantitatively.\n",
    "\n",
    "<b>Note</b>: A more detailed version of the approach described here can be found in the [Segmentation Evaluation notebook](34_Segmentation_Evaluation.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Transfer the segmentation via the estimated transformation. Use Nearest Neighbor interpolation to retain the labels.\n",
    "transformed_labels = sitk.Resample(masks[moving_image_index],\n",
    "                                   images[fixed_image_index],\n",
    "                                   tx, \n",
    "                                   sitk.sitkNearestNeighbor,\n",
    "                                   0.0, \n",
    "                                   masks[moving_image_index].GetPixelID())\n",
    "\n",
    "segmentations_before_and_after = [masks[moving_image_index], transformed_labels]\n",
    "interact(display_coronal_with_label_maps_overlay, coronal_slice = (0, images[0].GetSize()[1]-1),\n",
    "         mask_index=(0,len(segmentations_before_and_after)-1),\n",
    "         image = fixed(images[fixed_image_index]), masks = fixed(segmentations_before_and_after), \n",
    "         label=fixed(lung_label), window_min = fixed(-1024), window_max=fixed(976));\n",
    "\n",
    "# Compute the Dice coefficient and Hausdorff distance between the segmentations before, and after registration.\n",
    "ground_truth = masks[fixed_image_index] == lung_label\n",
    "before_registration = masks[moving_image_index] == lung_label\n",
    "after_registration = transformed_labels == lung_label\n",
    "\n",
    "label_overlap_measures_filter = sitk.LabelOverlapMeasuresImageFilter()\n",
    "label_overlap_measures_filter.Execute(ground_truth, before_registration)\n",
    "print(\"Dice coefficient before registration: {:.2f}\".format(label_overlap_measures_filter.GetDiceCoefficient()))\n",
    "label_overlap_measures_filter.Execute(ground_truth, after_registration)\n",
    "print(\"Dice coefficient after registration: {:.2f}\".format(label_overlap_measures_filter.GetDiceCoefficient()))\n",
    "\n",
    "hausdorff_distance_image_filter = sitk.HausdorffDistanceImageFilter()\n",
    "hausdorff_distance_image_filter.Execute(ground_truth, before_registration)\n",
    "print(\"Hausdorff distance before registration: {:.2f}\".format(hausdorff_distance_image_filter.GetHausdorffDistance()))\n",
    "hausdorff_distance_image_filter.Execute(ground_truth, after_registration)\n",
    "print(\"Hausdorff distance after registration: {:.2f}\".format(hausdorff_distance_image_filter.GetHausdorffDistance()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "## Multi-resolution control point grid\n",
    "\n",
    "In the example above we used the standard image registration framework. This implies the same transformation model at all image resolutions. For global transformations (e.g. rigid, affine...) the number of transformation parameters has no relationship to the changing resolution. For the BSpline transformation we can potentially use fewer control points for images with lower frequencies, higher levels of the image pyramid, increasing the number of control points as we go down the pyramid. With the standard framework we use the same number of control points for all pyramid levels.\n",
    "\n",
    "To use a multi-resolution control point grid we have a specific initializer for the BSpline transformation, [SetInitialTransformAsBSpline](https://itk.org/SimpleITKDoxygen/html/classitk_1_1simple_1_1ImageRegistrationMethod.html#ada2d2f32fda2b2aa70f8ddfe40bfd74d).\n",
    "\n",
    "The following code solves the same registration task as above, just with a multi-resolution control point grid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def bspline_intra_modal_registration2(fixed_image, moving_image, fixed_image_mask=None, fixed_points=None, moving_points=None):\n",
    "\n",
    "    registration_method = sitk.ImageRegistrationMethod()\n",
    "    \n",
    "    # Determine the number of BSpline control points using the physical spacing we \n",
    "    # want for the finest resolution control grid. \n",
    "    grid_physical_spacing = [50.0, 50.0, 50.0] # A control point every 50mm\n",
    "    image_physical_size = [size*spacing for size,spacing in zip(fixed_image.GetSize(), fixed_image.GetSpacing())]\n",
    "    mesh_size = [int(image_size/grid_spacing + 0.5) \\\n",
    "                 for image_size,grid_spacing in zip(image_physical_size,grid_physical_spacing)]\n",
    "\n",
    "    # The starting mesh size will be 1/4 of the original, it will be refined by \n",
    "    # the multi-resolution framework.\n",
    "    mesh_size = [int(sz/4 + 0.5) for sz in mesh_size]\n",
    "    \n",
    "    initial_transform = sitk.BSplineTransformInitializer(image1 = fixed_image, \n",
    "                                                         transformDomainMeshSize = mesh_size, order=3)    \n",
    "    # Instead of the standard SetInitialTransform we use the BSpline specific method which also\n",
    "    # accepts the scaleFactors parameter to refine the BSpline mesh. In this case we start with \n",
    "    # the given mesh_size at the highest pyramid level then we double it in the next lower level and\n",
    "    # in the full resolution image we use a mesh that is four times the original size.\n",
    "    registration_method.SetInitialTransformAsBSpline(initial_transform,\n",
    "                                                     inPlace=True,\n",
    "                                                     scaleFactors=[1,2,4])\n",
    "    registration_method.SetMetricAsMeanSquares()\n",
    "    # Settings for metric sampling, usage of a mask is optional. When given a mask the sample points will be \n",
    "    # generated inside that region. Also, this implicitly speeds things up as the mask is smaller than the\n",
    "    # whole image.\n",
    "    registration_method.SetMetricSamplingStrategy(registration_method.RANDOM)\n",
    "    registration_method.SetMetricSamplingPercentage(0.01)\n",
    "    if fixed_image_mask:\n",
    "        registration_method.SetMetricFixedMask(fixed_image_mask)\n",
    "    \n",
    "    # Multi-resolution framework.            \n",
    "    registration_method.SetShrinkFactorsPerLevel(shrinkFactors = [4,2,1])\n",
    "    registration_method.SetSmoothingSigmasPerLevel(smoothingSigmas=[2,1,0])\n",
    "    registration_method.SmoothingSigmasAreSpecifiedInPhysicalUnitsOn()\n",
    "\n",
    "    registration_method.SetInterpolator(sitk.sitkLinear)\n",
    "    # Use the LBFGS2 instead of LBFGS. The latter cannot adapt to the changing control grid resolution.\n",
    "    registration_method.SetOptimizerAsLBFGS2(solutionAccuracy=1e-2, numberOfIterations=100, deltaConvergenceTolerance=0.01)\n",
    "\n",
    "    # If corresponding points in the fixed and moving image are given then we display the similarity metric\n",
    "    # and the TRE during the registration.\n",
    "    if fixed_points and moving_points:\n",
    "        registration_method.AddCommand(sitk.sitkStartEvent, rc.metric_and_reference_start_plot)\n",
    "        registration_method.AddCommand(sitk.sitkEndEvent, rc.metric_and_reference_end_plot)\n",
    "        registration_method.AddCommand(sitk.sitkIterationEvent, lambda: rc.metric_and_reference_plot_values(registration_method, fixed_points, moving_points))\n",
    "    \n",
    "    return registration_method.Execute(fixed_image, moving_image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#%%timeit -r1 -n1\n",
    "\n",
    "# Select the fixed and moving images, valid entries are in [0,9].\n",
    "fixed_image_index = 0\n",
    "moving_image_index = 7\n",
    "\n",
    "\n",
    "tx = bspline_intra_modal_registration2(fixed_image = images[fixed_image_index], \n",
    "                                      moving_image = images[moving_image_index],\n",
    "                                      fixed_image_mask = (masks[fixed_image_index] == lung_label),\n",
    "                                      fixed_points = points[fixed_image_index], \n",
    "                                      moving_points = points[moving_image_index]\n",
    "                                     )\n",
    "initial_errors_mean, initial_errors_std, _, initial_errors_max, initial_errors = ru.registration_errors(sitk.Euler3DTransform(), points[fixed_image_index], points[moving_image_index])\n",
    "final_errors_mean, final_errors_std, _, final_errors_max, final_errors = ru.registration_errors(tx, points[fixed_image_index], points[moving_image_index])\n",
    "\n",
    "plt.hist(initial_errors, bins=20, alpha=0.5, label='before registration', color='blue')\n",
    "plt.hist(final_errors, bins=20, alpha=0.5, label='after registration', color='green')\n",
    "plt.legend()\n",
    "plt.title('TRE histogram');\n",
    "print('Initial alignment errors in millimeters, mean(std): {:.2f}({:.2f}), max: {:.2f}'.format(initial_errors_mean, initial_errors_std, initial_errors_max))\n",
    "print('Final alignment errors in millimeters, mean(std): {:.2f}({:.2f}), max: {:.2f}'.format(final_errors_mean, final_errors_std, final_errors_max))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
